import cv2
import numpy as np
from PIL import Image

def calculate_sharpness(image):
    """
    计算图像清晰度（基于拉普拉斯方差）
    参数:
        image: 输入图像(numpy数组)
    返回:
        清晰度分数(越高表示越清晰)
    """
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    return cv2.Laplacian(gray, cv2.CV_64F).var()

def calculate_brightness(image):
    """
    计算图像平均亮度
    参数:
        image: 输入图像
    返回:
        亮度值(0-255)
    """
    hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
    return np.mean(hsv[:,:,2])

def calculate_contrast(image):
    """
    计算图像对比度
    参数:
        image: 输入图像
    返回:
        对比度分数
    """
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    return gray.std()

def check_text_legibility(image, min_sharpness=50, min_brightness=80, max_brightness=200, min_contrast=40):
    """
    综合判断图像中的文字是否清晰可读
    参数:
        image: 输入图像(numpy数组)
        min_sharpness: 最小清晰度阈值
        min_brightness: 最小亮度阈值
        max_brightness: 最大亮度阈值
        min_contrast: 最小对比度阈值
    返回:
        dict: 包含各项指标和综合判断结果
    """
    # 计算各项指标
    sharpness = calculate_sharpness(image)
    brightness = calculate_brightness(image)
    contrast = calculate_contrast(image)
    
    # 判断各项指标是否达标
    is_sharp = sharpness >= min_sharpness
    is_brightness_ok = min_brightness <= brightness <= max_brightness
    is_contrast_ok = contrast >= min_contrast
    
    # 综合判断
    is_legible = is_sharp and is_brightness_ok and is_contrast_ok
    
    return {
        'sharpness': sharpness,
        'brightness': brightness,
        'contrast': contrast,
        'is_sharp': is_sharp,
        'is_brightness_ok': is_brightness_ok,
        'is_contrast_ok': is_contrast_ok,
        'is_legible': is_legible,
        'suggestions': get_suggestions(is_sharp, is_brightness_ok, is_contrast_ok)
    }

def get_suggestions(is_sharp, is_brightness_ok, is_contrast_ok):
    """根据检测结果给出改进建议"""
    suggestions = []
    if not is_sharp:
        suggestions.append("图像模糊，建议重新拍摄或使用锐化处理")
    if not is_brightness_ok:
        suggestions.append("亮度不合适，建议调整照明条件")
    if not is_contrast_ok:
        suggestions.append("对比度不足，建议增强文字与背景的对比")
    return suggestions if suggestions else ["图像质量良好，适合OCR识别"]

def preprocess_if_needed(image_path, output_path=None):
    """
    根据需要自动预处理图像
    参数:
        image_path: 输入图像路径
        output_path: 输出图像路径(可选)
    返回:
        预处理后的图像和检测报告
    """
    # 读取图像
    image = cv2.imread(image_path)
    if image is None:
        raise ValueError("无法读取图像，请检查路径")
    
    # 检测图像质量
    report = check_text_legibility(image)
    
    # 如果需要预处理且质量不达标
    if not report['is_legible']:
        # 自动预处理流程
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        
        # 亮度调整
        if not report['is_brightness_ok']:
            alpha = 1.5 if report['brightness'] < 80 else 0.8
            gray = cv2.convertScaleAbs(gray, alpha=alpha, beta=0)
        
        # 对比度增强
        if not report['is_contrast_ok']:
            clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
            gray = clahe.apply(gray)
        
        # 锐化处理
        if not report['is_sharp']:
            kernel = np.array([[-1,-1,-1], [-1,9,-1], [-1,-1,-1]])
            gray = cv2.filter2D(gray, -1, kernel)
        
        # 二值化
        _, processed = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
        
        # 保存结果
        if output_path:
            cv2.imwrite(output_path, processed)
        
        return processed, report
    
    # 如果质量达标，直接返回原图
    if output_path:
        cv2.imwrite(output_path, image)
    
    return image, report

# 使用示例
if __name__ == "__main__":
    image_path = "D:/ocr/sample/6cut1.jpg"  # 替换为你的图片路径
    
    # 方案1: 只检测不处理
    image = cv2.imread(image_path)
    report = check_text_legibility(image)
    
    print("\n图像质量检测报告:")
    print(f"清晰度: {report['sharpness']:.2f} ({'达标' if report['is_sharp'] else '不足'})")
    print(f"亮度: {report['brightness']:.2f} ({'合适' if report['is_brightness_ok'] else '不合适'})")
    print(f"对比度: {report['contrast']:.2f} ({'足够' if report['is_contrast_ok'] else '不足'})")
    print(f"综合判断: {'适合OCR识别' if report['is_legible'] else '可能需要预处理'}")
    print("改进建议:")
    for suggestion in report['suggestions']:
        print(f"- {suggestion}")
    
    # 方案2: 自动检测并预处理
    processed_img, process_report = preprocess_if_needed(image_path, "preprocessed.jpg")
    
    # 显示结果
    cv2.imshow("Original Image", cv2.resize(image, (800, 600)))
    if not process_report['is_legible']:
        cv2.imshow("Processed Image", cv2.resize(processed_img, (800, 600)))
    cv2.waitKey(0)
    cv2.destroyAllWindows()